# SPDX-FileCopyrightText: Copyright (c) 2018-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import re
import sys
import textwrap

import toml

VSCODE_SETTINGS_TEMPLATE = """
{
    "editor.rulers": [120],

    // Enables python language server (seems to work slightly better than jedi)
    "python.languageServer": "Pylance",
    "python.jediEnabled": false,

    // Those paths are automatically filled by isaacsim (see: 'python -m isaacsim --help')
    "python.defaultInterpreterPath": "PYTHON.DEFAULTINTERPRETERPATH",
    "python.analysis.extraPaths": [
        PYTHON.ANALYSIS.EXTRAPATHS
    ],

    // Use "black" as a formatter
    "python.formatting.provider": "black",
    "python.formatting.blackArgs": ["--line-length", "120"],

    // Use "flake8" for linting
    "python.linting.pylintEnabled": false,
    "python.linting.flake8Enabled": true,
}
"""


def generate_vscode_settings():
    def _mock_python_modules(ext_path: str, ext_name: str) -> None:
        # parse config/extension.toml
        cprint(f"  |-- Parsing extension config ({ext_path})")
        config_path = os.path.join(ext_path, "config", "extension.toml")
        try:
            config = toml.load(config_path)
        except Exception as e:
            cprint(f"  |     |-- [Warning] {e}")
            return
        # get python modules
        for item in config.get("python", {}).get("module", []):
            if list(item.keys()) == ["name"]:
                # skip tests
                if item.get("name", "").endswith(".tests"):
                    continue
                # mock __init__.py for each submodule (if not exists)
                submodule_path = ext_path
                for submodule in item.get("name", "").split("."):
                    init_path = os.path.join(submodule_path, "__init__.py")
                    if not os.path.isfile(init_path):
                        try:
                            cprint(f"  |-- Mocking {init_path}")
                            with open(init_path, "w") as f:
                                f.write("# Generated by 'isaacsim' package")
                        except Exception as e:
                            cprint(f"  |     |-- [Warning] {e}")
                            continue
                    submodule_path = os.path.join(submodule_path, submodule)

    def _get_paths(base_path: str, mock_python_modules: bool = False) -> list[str]:
        paths = []
        if os.path.isdir(base_path):
            for folder in os.listdir(base_path):
                folder_path = os.path.join(base_path, folder)
                if os.path.isdir(folder_path):
                    paths.append(folder_path)
                    cprint(f"Registering extension: {folder_path}")
                    if mock_python_modules:
                        _mock_python_modules(folder_path, re.split(r"-\d+", folder)[0])
        return paths

    # isaac-sim path
    isaacsim_path = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))

    # kit path (internal kernel)
    if os.path.isdir(os.path.join(isaacsim_path, "kit", "extscore")):
        kit_path = os.path.join(isaacsim_path, "kit")
        if kit_path not in sys.path:
            sys.path.append(kit_path)
        try:
            import kit_app  # importing 'kit_app' will bootstrap kernel

        except Exception as e:
            sys.exit(f"Unable to bootstrap inner kit kernel: {e}")
    # kit path (omniverse-kit kernel package)
    else:
        try:
            import omni.kit_app  # importing 'kit_app' will bootstrap kernel

            kit_path = os.path.dirname(os.path.abspath(os.path.realpath(omni.kit_app.__file__)))
        except ModuleNotFoundError:
            sys.exit("Unable to find 'omniverse-kit' package")

    cwd = os.getcwd()
    vscode_settings_path = os.path.join(cwd, ".vscode", "settings.json")
    # check if .vscode/settings.json exists
    if os.path.exists(vscode_settings_path):
        print(f"VS Code settings already exists: {vscode_settings_path}")
        if input("Overwrite? (y/N): ").lower() not in ["y", "yes"]:
            print("Cancelled: VS Code settings not overwritten")
            return

    # get extensions paths
    extensions_paths = []
    # - omniverse-kit
    folder_path = os.path.join(kit_path, "kernel", "py")
    if os.path.isdir(folder_path):
        extensions_paths.append(folder_path)
    for folder in ["exts", "extscore"]:
        extensions_paths.extend(_get_paths(os.path.join(kit_path, folder), mock_python_modules=True))
    # - isaacsim
    for folder in ["exts", "extscache", "extsDeprecated", "extsUser"]:
        extensions_paths.extend(_get_paths(os.path.join(isaacsim_path, folder), mock_python_modules=True))

    # update 'python.defaultInterpreterPath'
    template = VSCODE_SETTINGS_TEMPLATE[:]
    template = template.replace("PYTHON.DEFAULTINTERPRETERPATH", sys.executable)

    # update 'python.analysis.extraPaths'
    content = "\n".join([f'"{path}",' for path in extensions_paths])
    content = textwrap.indent(content, prefix=" " * 8)[8:]
    template = template.replace("PYTHON.ANALYSIS.EXTRAPATHS", content)

    # create .vscode/settings.json
    os.makedirs(os.path.join(cwd, ".vscode"), exist_ok=True)
    with open(vscode_settings_path, "w") as f:
        f.write(template)
    print("VS Code settings generated at", vscode_settings_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--generate-vscode-settings", default=False, action="store_true", help="Generate VS Code settings."
    )
    parser.add_argument("--verbose", default=False, action="store_true", help="Verbose output.")
    args, _ = parser.parse_known_args()

    cprint = print if args.verbose else lambda *args, **kwargs: None

    if args.generate_vscode_settings:
        generate_vscode_settings()
